#include "../includes/TaskConsumer.h"
#include "../includes/TaskManager.h"
#include <iostream>
#include <thread>
#include <atomic>
#include <chrono>
#include <string>
#include <vector>
#include <sstream>
#include <memory>
#include <fstream>      // 用于文件日志
#include <iomanip>      // 用于时间格式化
#include <ctime>        // 用于时间戳
#include <unistd.h>     // 用于 getpid()
#include <sys/types.h>  // 用于 pid_t
#include <cstdio>       // 用于 popen, pclose, fgets
#include <array>        // 用于 exec 辅助函数
#include <algorithm>    // 用于 std::remove

/**
 * @brief 执行一个 shell 命令并返回其标准输出
 * @param cmd 要执行的命令
 * @return 命令的标准输出字符串
 */
std::string exec(const char* cmd) {
    std::array<char, 128> buffer;
    std::string result;
    // 使用 std::unique_ptr 自动管理 FILE* 指针
    std::unique_ptr<FILE, decltype(&pclose)> pipe(popen(cmd, "r"), pclose);
    if (!pipe) {
        return "popen() failed!";
    }
    while (fgets(buffer.data(), buffer.size(), pipe.get()) != nullptr) {
        result += buffer.data();
    }
    return result;
}

/**
 * @brief 构建一个任务请求字符串
 * @param task_id 任务ID
 * @param task_type 任务类型 (0-4)
 * @param filename (可选) 相关的文件名
 * @param content (可选) 写入的内容
 * @return 序列化后的任务字符串
 */
std::string build_task_request(int32_t task_id, char task_type, const std::string& filename = "", const std::string& content = "") {
    std::string task_request;
    int32_t id_int = task_id;
    char type_byte = task_type;

    task_request.append(reinterpret_cast<const char*>(&id_int), 4);
    task_request.append(&type_byte, 1);

    switch (task_type) {
        case 0: { // 写文件: +1字节fn_len + filename + 4字节content_len + content
            if (filename.length() > 128 || content.length() > 1024) return "";
            char filename_len_byte = static_cast<char>(filename.length());
            int32_t content_len_int = content.length();
            task_request.append(&filename_len_byte, 1);
            task_request.append(filename);
            task_request.append(reinterpret_cast<const char*>(&content_len_int), 4);
            task_request.append(content);
            break;
        }
        case 1: // 读文件: +1字节fn_len + filename
        case 2: // 创建文件: +1字节fn_len + filename
        case 3: // 删除文件: +1字节fn_len + filename
        {
            if (filename.length() > 128) return "";
            char filename_len_byte = static_cast<char>(filename.length());
            task_request.append(&filename_len_byte, 1);
            task_request.append(filename);
            break;
        }
        case 4: // 查看目录: 无额外内容
            break;
        default:
            return ""; // 未知类型
    }
    return task_request;
}


int main() {
    // --- 配置测试时长 ---
    // 在这里修改总测试时长 (小时)
    double test_duration_hours = 12.0; 
    // double test_duration_hours = 0.01; // 用于快速调试 (约36秒)
    
    auto test_duration_seconds = static_cast<long long>(test_duration_hours * 3600);

    // --- 初始化 ---
    TaskManager& task_manager = TaskManager::getInstance();
    TaskConsumer& task_consumer = TaskConsumer::getInstance(); 

    std::atomic<bool> running(true); // 控制所有线程的运行状态
    std::atomic<int> next_task_id(1); // 用于生成唯一的任务ID

    // --- 准备监控日志 ---
    std::ofstream metrics_log("../logs/stress_test_metrics.txt");
    if (!metrics_log.is_open()) {
        std::cerr << "严重错误: 无法打开 'stress_test_metrics.txt' 日志文件!" << std::endl;
        return 1;
    }
    metrics_log << "Timestamp, CPU (%), Memory (%), Disk Usage (Current Dir)" << std::endl;

    pid_t pid = getpid();
    std::string ps_command = "ps -p " + std::to_string(pid) + " -o %cpu,%mem --no-headers";
    std::string du_command = "du -sh . | awk '{print $1}'"; // 仅获取目录大小


    // --- 线程1: 生产者 - 持续按顺序生成任务 ---
    std::thread producer_thread([&task_manager, &running, &next_task_id]() {
        std::string base_filename = "stress_file_";
        
        while (running) {
            // 为一个文件操作序列保留5个连续的ID
            int file_op_id = next_task_id.fetch_add(5); 
            std::string filename = base_filename + std::to_string(file_op_id) + ".txt";

            // --- 1. 创建 (Touch) Task (Type 2) ---
            int task_id_touch = file_op_id;
            std::string touch_req = build_task_request(task_id_touch, 2, filename);
            if (!touch_req.empty() && running) {
                task_manager.add_task(touch_req);
                std::cout << "Producer: Queued Touch (ID: " << task_id_touch << ", File: " << filename << ")" << std::endl;
            }
            if (!running) break;
            std::this_thread::sleep_for(std::chrono::milliseconds(10)); // 任务间短暂间隔

            // --- 2. 写 (Write) Task (Type 0) ---
            int task_id_write = file_op_id + 1;
            std::string content = "Stress test content for file " + filename + " at " + std::to_string(std::time(nullptr)) + "\n";
            std::string write_req = build_task_request(task_id_write, 0, filename, content);
            if (!write_req.empty() && running) {
                task_manager.add_task(write_req);
                std::cout << "Producer: Queued Write (ID: " << task_id_write << ", File: " << filename << ")" << std::endl;
            }
            if (!running) break;
            std::this_thread::sleep_for(std::chrono::milliseconds(10));

            // --- 3. 读 (Cat) Task (Type 1) ---
            int task_id_read = file_op_id + 2;
            std::string read_req = build_task_request(task_id_read, 1, filename);
            if (!read_req.empty() && running) {
                task_manager.add_task(read_req);
                std::cout << "Producer: Queued Read (ID: " << task_id_read << ", File: " << filename << ")" << std::endl;
            }
            if (!running) break;
            std::this_thread::sleep_for(std::chrono::milliseconds(10));

            // --- 4. 删 (Rm) Task (Type 3) ---
            int task_id_rm = file_op_id + 3;
            std::string rm_req = build_task_request(task_id_rm, 3, filename);
            if (!rm_req.empty() && running) {
                task_manager.add_task(rm_req);
                std::cout << "Producer: Queued Rm (ID: " << task_id_rm << ", File: " << filename << ")" << std::endl;
            }
            if (!running) break;
            std::this_thread::sleep_for(std::chrono::milliseconds(10));

            // --- 5. 查 (Ls) Task (Type 4) ---
            int task_id_ls = file_op_id + 4;
            std::string ls_req = build_task_request(task_id_ls, 4);
            if (!ls_req.empty() && running) {
                task_manager.add_task(ls_req);
                std::cout << "Producer: Queued Ls (ID: " << task_id_ls << ")" << std::endl;
            }

            // --- 在下一个文件操作循环前休眠 ---
            // 调整此值可控制文件操作的总体速率
            if (running) {
                std::this_thread::sleep_for(std::chrono::milliseconds(200)); 
            }
        }
        std::cout << "Producer thread stopping." << std::endl;
    });

    // --- 线程2: 分发者 - (逻辑不变) ---
    std::thread dispatcher_thread([&task_manager, &task_consumer, &running]() {
        while (running) {
            if (task_consumer.idle_threads() > 0 && !task_manager.task_empty()) {
                if (!task_consumer.add_task()) { 
                    std::cerr << "Dispatcher: Failed to add task to TaskConsumer." << std::endl;
                } else {
                    // std::cout << "Dispatcher: Submitted task to TaskConsumer." << std::endl; // 日志太多，注释掉
                }
            } else {
                if (running) {
                    std::this_thread::sleep_for(std::chrono::milliseconds(10));
                }
            }
        }
        std::cout << "Dispatcher thread stopping." << std::endl;
    });

    // --- 线程3: 消费者 - (逻辑不变) ---
    std::thread consumer_thread([&task_manager, &running]() {
        while (running) {
            if (!task_manager.resp_empty()) { 
                std::string resp_str;
                task_manager.get_resp(resp_str); 

                if (!resp_str.empty()) {
                    // --- 解析响应内容 (假设内容是 [任务ID(4字节) + 类型ID(1字节) + 状态码(1字节) + ...]) ---
                    if (resp_str.length() < 6) { 
                        std::cerr << "Consumer: Received response content is too short." << std::endl;
                        continue;
                    }

                    size_t offset = 0;
                    int task_id = *(reinterpret_cast<const int*>(resp_str.data() + offset));
                    offset += 4;
                    char type_id = resp_str[offset];
                    offset += 1;
                    char status_code = resp_str[offset];
                    offset += 1;

                    std::cout << "Consumer: Received Response - Task ID: " << task_id << ", Type: " << static_cast<int>(type_id) 
                              << ", Status Code: " << static_cast<int>(status_code);

                    std::string response_details = "";

                    switch (type_id) {
                        case 0: // 写文件响应: +4字节bytes_written
                            if (resp_str.length() >= offset + 4) {
                                int bytes_written = *(reinterpret_cast<const int*>(resp_str.data() + offset));
                                response_details = ", Bytes Written: " + std::to_string(bytes_written);
                            } else {
                                std::cerr << " [ERROR] Write response content is truncated." << std::endl;
                            }
                            break;
                        case 1: // 读文件响应: +4字节content_length + 变长content
                            if (resp_str.length() >= offset + 4) {
                                int32_t content_len = *(reinterpret_cast<const int*>(resp_str.data() + offset));
                                if (resp_str.length() >= offset + 4 + content_len) {
                                    std::string content_str = resp_str.substr(offset + 4, content_len);
                                    // 为避免刷屏，只显示内容长度
                                    response_details = ", Content Length: " + std::to_string(content_len); // ", Content: " + content_str;
                                } else {
                                    std::cerr << " [ERROR] Cat response content is truncated." << std::endl;
                                }
                            } else {
                                std::cerr << " [ERROR] Cat response content is truncated." << std::endl;
                            }
                            break;
                        case 2: // 创建文件响应: 无额外内容
                        case 3: // 删除文件响应: 无额外内容
                            break;
                        case 4: // 查看目录响应: +1字节file_count + 变长file_entries
                            if (resp_str.length() >= offset + 1) {
                                char file_count = resp_str[offset];
                                offset += 1;
                                std::string entries_str = resp_str.substr(offset);
                                // 为避免刷屏，只显示文件计数
                                response_details = ", File Count: " + std::to_string(static_cast<int>(file_count));
                            } else {
                                std::cerr << " [ERROR] Ls response content is truncated." << std::endl;
                            }
                            break;
                        default:
                            std::cerr << " [ERROR] Unknown response type ID: " << static_cast<int>(type_id) << std::endl;
                            break;
                    }
                    std::cout << response_details << std::endl;
                }
            } else {
                if (running) {
                    std::this_thread::sleep_for(std::chrono::milliseconds(10));
                }
            }
        }
        std::cout << "Consumer thread stopping." << std::endl;
    });

    // --- 线程4: 监控者 - 持续记录系统指标 ---
    std::thread monitor_thread([&running, &metrics_log, ps_command, du_command]() {
        while (running) {
            // 获取当前时间戳
            auto now = std::chrono::system_clock::now();
            std::time_t now_c = std::chrono::system_clock::to_time_t(now);
            std::tm now_tm = *std::localtime(&now_c);
            char time_buf[80];
            std::strftime(time_buf, sizeof(time_buf), "%Y-%m-%d %H:%M:%S", &now_tm);

            // 执行命令
            std::string ps_output = exec(ps_command.c_str());
            std::string du_output = exec(du_command.c_str());
            
            // 清理 ps 输出 (例如: "  0.1  0.5")
            std::stringstream ss_ps(ps_output);
            std::string cpu, mem;
            ss_ps >> cpu >> mem;
            
            // 清理 du 输出 (例如: "1.2G\n")
            std::string disk_usage = du_output;
            disk_usage.erase(std::remove(disk_usage.begin(), disk_usage.end(), '\n'), disk_usage.end());

            // 打印到控制台
            std::cout << "[Monitor] Time: " << time_buf << ", CPU: " << cpu << "%, Mem: " << mem << "%, Disk: " << disk_usage << std::endl;
            
            // 写入日志文件
            metrics_log << time_buf << ", " << cpu << ", " << mem << ", " << disk_usage << std::endl;

            // 每 5 秒记录一次
            for(int i = 0; i < 60 && running; ++i) {
                std::this_thread::sleep_for(std::chrono::seconds(1));
            }
        }
        std::cout << "Monitor thread stopping." << std::endl;
    });


    // --- 主线程等待测试结束 ---
    
    std::cout << "Stress test started. Duration: " << test_duration_hours << " hours (" << test_duration_seconds << " seconds)." << std::endl;
    std::cout << "Monitoring data will be saved to 'stress_test_metrics.txt'" << std::endl;
    
    // 主线程等待测试时长
    std::this_thread::sleep_for(std::chrono::seconds(test_duration_seconds));

    std::cout << "Test duration elapsed. Stopping all threads..." << std::endl;
    running = false; // 向所有线程发送停止信号

    // 等待所有线程结束
    if (producer_thread.joinable()) {
        producer_thread.join();
    }
    if (dispatcher_thread.joinable()) {
        dispatcher_thread.join();
    }
    if (consumer_thread.joinable()) {
        consumer_thread.join();
    }
    if (monitor_thread.joinable()) {
        monitor_thread.join();
    }

    metrics_log.close(); // 关闭日志文件
    std::cout << "All threads finished. Stress test completed." << std::endl;
    return 0;
}
